{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vPzPjEAanpOL"
      },
      "outputs": [],
      "source": [
        "from transformers import ViTForImageClassification, ViTImageProcessor\n",
        "\n",
        "model = ViTForImageClassification.from_pretrained(\"google/vit-base-patch16-224\")\n",
        "processor = ViTImageProcessor.from_pretrained(\"google/vit-base-patch16-224\")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from PIL import Image\n",
        "import requests\n",
        "url = 'http://images.cocodataset.org/val2017/000000439715.jpg'\n",
        "image = Image.open(requests.get(url, stream=True).raw)"
      ],
      "metadata": {
        "id": "kryX6kS9nrlT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "inputs = processor(images=image, return_tensors=\"pt\")\n",
        "outputs = model(inputs.pixel_values)\n",
        "prediction = outputs.logits.argmax(-1)\n",
        "print(\"Class label: \", model.config.id2label[prediction.item()])"
      ],
      "metadata": {
        "id": "hyeElfX3nt4u"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}